import cv2
import detect
import os, time
from tqdm import tqdm 

# 图像获取
gray_folder = "E:\Python\PythonProject\detect\dataset\gray"
depth_folder = "E:\Python\PythonProject\detect\dataset\depth"
# 获取文件夹中的所有文件名
gray_files = [f for f in os.listdir(gray_folder) if f.endswith('.jpg')]
depth_files = [f for f in os.listdir(depth_folder) if f.endswith('.jpg')]
img_nameList = [
    "1226.jpg","3383.jpg","1114.jpg","1103.jpg","0152.jpg","0149.jpg","0160.jpg","3397.jpg","0058.jpg","0000.jpg", "0040.jpg", "0057.jpg", "0131.jpg", "1100.jpg"
]
i = 0
detector = detect.RebarDetector(model_path="model.yml")
# 确保两个文件夹中的文件数量相同
if len(gray_files) != len(depth_files):
    print("灰度图和深度图文件数量不匹配")
else:
    total_time = 0.0
    for img_name in tqdm(gray_files, desc="处理进度", unit="图像"):
        start_time = time.time() 
        # print(img_name)  # 保留原有打印（可选）
        # 读取灰度图和深度图
        ir_img_path = os.path.join(gray_folder, img_name)
        depth_img_path = os.path.join(depth_folder, img_name)
        ir_img = cv2.imread(ir_img_path, 0)
        depth_img = cv2.imread(depth_img_path, cv2.COLOR_BGR2RGB)
        
        # 处理图像
        result = detector.process_image(ir_img, depth_img)
        
        if result is not None:
            # 保存结果图像
            result_path = os.path.join("results", img_name)
            cv2.imwrite(result_path, result)
            # print(f"结果已保存到: {result_path}")  # 保留原有打印（可选）
            
        # 计算耗时
        end_time = time.time()
        iter_time = end_time - start_time
        total_time += iter_time
        
        # 在进度条后显示单张处理时间（可选）
        # tqdm.write(f"处理 {img_name} 耗时: {iter_time:.2f} 秒")  

    # 最终统计
    print(f"\n平均处理时间: {total_time / len(gray_files):.2f} 秒")

        